import gradio as gr
import random
import time
from langchain_ollama import ChatOllama

llm = ChatOllama(base_url="127.0.0.1:11434", model="qwen2.5:14b", num_ctx=4096, temperature=0.7, streaming=True)


with gr.Blocks() as demo:
    chatbot = gr.Chatbot(type="messages")
    msg = gr.Textbox()
    clear = gr.Button("Clear")

    def user(user_message, history: list):
        return "", history + [{"role": "user", "content": user_message}]

    def bot(history: list):
        # bot_message = random.choice(["How are you?", "I love you", "I'm very hungry"])
        # bot_message = llm.invoke(history)
        bot_message = llm.stream(history)
        history.append({"role": "assistant", "content": ""})
        for chunk in bot_message:
            content = chunk.content
            history[-1]['content'] += content
            yield history
        # for character in bot_message.content:
        #     history[-1]['content'] += character
        #     time.sleep(0.005)
        #     yield history
        # full_response = ""
        # for response_chunk in llm.stream(history):
        #     content = response_chunk.content
        #     full_response += content
        #     yield full_response

    msg.submit(user, [msg, chatbot], [msg, chatbot], queue=False).then(
        bot, chatbot, chatbot
    )
    clear.click(lambda: None, None, chatbot, queue=False)

demo.launch()